# src/calphad/calphadpy.py

import os

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from xarray import DataArray, Dataset
from pycalphad import Database, equilibrium, variables as v

from src.calphad.base import CalphadBase


def squeeze_array(data: DataArray) -> DataArray:
    return data.squeeze(drop=True)


class CalphadPy(CalphadBase):

    def __init__(self, tdb_file: str) -> None:
        super().__init__(tdb_file)
        self.dbf = Database(tdb_file)

    def get_species(self) -> list[str]:
        return [sp.name for sp in self.dbf.species]

    def get_phases(self) -> list[str]:
        return list(self.dbf.phases.keys())

    def calculate_equilibrium(self, comps_fracts, phases, temps, pressure=1e5) -> Dataset:
        comps_fracts = [(c.upper(), f) for c, f in sorted(comps_fracts.items())]

        # Set components and phases
        comps, fracts = zip(*comps_fracts)
        phases = [phase for phase in phases if phase in self.get_phases()]

        # Set temperatures
        if isinstance(temps, tuple) and len(temps) == 3:
            temps = np.arange(*temps).tolist() + [temps[1]]

        # Set conditions
        comps_fracts.pop(0)
        mole_fracts = {v.X(c): f for c, f in comps_fracts}
        conditions = {v.P: pressure, v.T: temps, **mole_fracts}

        # Calculate
        result: Dataset = equilibrium(dbf=self.dbf, comps=comps, phases=phases, conditions=conditions)

        return result

    def plot_GM_T(self, outdir: str, data: Dataset) -> pd.DataFrame:
        if "GM" not in data.data_vars:
            raise ValueError("GM not found in data")

        # Squeeze array
        GM = squeeze_array(data["GM"])

        plt.figure(figsize=(10, 6))
        gm_vals = GM

        # Filter out NaNs
        valid = ~np.isnan(gm_vals)
        if not valid.any():
            raise ValueError("No valid data found")
        t_vals = GM["T"].values[valid]
        gm_vals = gm_vals.values[valid]

        # Plot
        plt.plot(t_vals, gm_vals, label="GM")
        plt.xlabel("Temperature (K)")
        plt.ylabel("Gibbs energy (J)")
        plt.title("GM vs T")
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(os.path.join(outdir, "GM-T.png"))

        # Save data
        return pd.DataFrame(data=gm_vals, index=t_vals, columns=["GM"])

    def plot_NP_T(self, outdir: str, data: Dataset) -> pd.DataFrame:
        if "NP" not in data.data_vars:
            raise ValueError("NP not found in data")
        if "Phase" not in data.data_vars:
            raise ValueError("Phase not found in data")

        # Squeeze arrays
        NP = squeeze_array(data["NP"])
        Phase = squeeze_array(data["Phase"])

        np_data = []
        plt.figure(figsize=(10, 6))
        for v in NP["vertex"].values:
            np_v = NP.sel(vertex=v)
            phase_v = Phase.sel(vertex=v)

            # Filter out NaNs
            valid = ~np.isnan(np_v)
            if not valid.any():
                print(f"Warning: No valid data for vertex {v}")
                continue

            t_vals = NP["T"].values[valid]
            np_vals = np_v.values[valid]
            phase_names = phase_v.values[valid]
            phase_label = [p for p in phase_names if p.strip()][0]

            # Plot
            label = f"{phase_label}#{v+1}"
            plt.plot(t_vals, np_vals, label=label)

            # Save data
            np_df = pd.DataFrame(data=np_vals, index=t_vals, columns=[label])
            np_data.append(np_df)

        plt.xlabel("Temperature (K)")
        plt.ylabel("Amount of phase")
        plt.title("NP vs T")
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(os.path.join(outdir, "NP-T.png"))
        plt.close()

        return pd.concat(np_data, axis=1)

    def plot_X_T(self, outdir: str, data: Dataset) -> pd.DataFrame:
        if "X" not in data.data_vars:
            raise ValueError("X not found in data")
        if "Phase" not in data.data_vars:
            raise ValueError("Phase not found in data")

        # Squeeze arrays
        X = squeeze_array(data["X"])
        Phase = squeeze_array(data["Phase"])

        x_data = []
        plt.figure(figsize=(12, 8))
        for v in X["vertex"].values:
            phase_v = Phase.sel(vertex=v)
            for c in X["component"].values:
                x_v_c = X.sel(vertex=v, component=c)

                # Filter out NaNs
                valid = ~np.isnan(x_v_c)
                if not valid.any():
                    print(f"Warning: No valid data for vertex {v} and component {c}")
                    continue

                t_vals = X["T"].values[valid]
                x_vals = x_v_c.values[valid]
                phase_names = phase_v.values[valid]
                phase_label = [p for p in phase_names if p.strip()][0]

                # Plot
                label = f"{phase_label}#{v+1},{c}"
                plt.plot(t_vals, x_vals, label=label)

                # Save data
                x_df = pd.DataFrame(data=x_vals, index=t_vals, columns=[label])
                x_data.append(x_df)

        plt.xlabel("Temperature (K)")
        plt.ylabel("Composistion of phase")
        plt.title("X vs T")
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(os.path.join(outdir, "X-T.png"))
        plt.close()

        return pd.concat(x_data, axis=1)

    def plot_Y_T(self, outdir: str, data: Dataset) -> pd.DataFrame:
        if "Y" not in data.data_vars:
            raise ValueError("Y not found in data")
        if "Phase" not in data.data_vars:
            raise ValueError("Phase not found in data")

        # Squeeze data
        Y = squeeze_array(data["Y"])
        Phase = squeeze_array(data["Phase"])
        comps = data["component"].values
        y_data = []
        plt.figure(figsize=(12, 8))
        for v in Y["vertex"].values:
            phase_v = Phase.sel(vertex=v)
            for d in Y["internal_dof"].values:
                y_v_d = Y.sel(vertex=v, internal_dof=d)

                # Filter out NaNs
                valid = ~np.isnan(y_v_d)
                if not valid.any():
                    print(f"Warning: No valid data for vertex {v} and internal dof {d}")
                    continue

                t_vals = Y["T"].values[valid]
                y_vals = y_v_d.values[valid]
                phase_names = phase_v.values[valid]
                phase_label = [p for p in phase_names if p.strip()][0]

                # Plot
                label = f"{phase_label}#{v+1},{comps[d % len(comps)]}#{d//len(comps)+1}"
                plt.plot(t_vals, y_vals, label=label)

                # Save data
                y_df = pd.DataFrame(data=y_vals, index=t_vals, columns=[label])
                y_data.append(y_df)

        plt.xlabel("Temperature (K)")
        plt.ylabel("Site Fraction")
        plt.title("Y vs T")
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(os.path.join(outdir, "Y-T.png"))
        plt.close()

        return pd.concat(y_data, axis=1)
